import os
import numpy as np
from sklearn.model_selection import train_test_split
import torch
import torchvision.transforms as tfs
import  torchvision.datasets as datas
from PIL import Image
DATA=r'E:/py-code/DL2/airplane-car-bird/' # 除了处理图片以外，不能放其他图片

# 读取图片和标签
def Read_Data(path):
    all_images = []
    all_labels = []
    train = datas.ImageFolder(root=path, transform=tfs.ToTensor())
    data = os.listdir(path+'data3')  # 所有图片名 列表
    for labels, (image, _) in enumerate(train):
        all_images.append(image.data.numpy())
        all_labels.append(int(data[labels][0]))

    return np.array(all_images), np.array(all_labels)
all_images,all_labels=Read_Data(DATA)

train_x,test_x,train_y,test_y=train_test_split(all_images,all_labels,test_size=0.2,shuffle=True)

train_x=torch.autograd.Variable(torch.Tensor(train_x)).cuda()
train_y=torch.autograd.Variable(torch.Tensor(train_y)).cuda()
test_x=torch.autograd.Variable(torch.Tensor(test_x)).cuda()
test_y=torch.autograd.Variable(torch.Tensor(test_y)).cuda()

class ConvBLK(torch.nn.Module):
    def __init__(self,ch_in,ch_out,stride=1):
        super(ConvBLK, self).__init__()
        self.conv1=torch.nn.Conv2d(ch_in,ch_out,kernel_size=3,stride=stride,padding=1)
        self.bn1=torch.nn.BatchNorm2d(ch_out)
        self.relu1=torch.nn.ReLU()
        
        self.conv2=torch.nn.Conv2d(ch_out,ch_out,kernel_size=3,stride=1,padding=1)
        self.bn2=torch.nn.BatchNorm2d(ch_out)

        if stride!=1:
            self.resiual=torch.nn.Sequential(torch.nn.Conv2d(ch_in,ch_out,kernel_size=1,stride=stride),
                                             torch.nn.BatchNorm2d(ch_out))
        else:
            self.resiual=lambda x:x
    def forward(self,x):
        out=self.conv1(x)
        out=self.bn1(out)
        out=self.relu1(out)
        out=self.conv2(out)
        out=self.bn2(out)
        r=self.resiual(x)

        output=out+r
        return output

class ResNet(torch.nn.Module):
    def __init__(self,num_blocks,nb_class):
        super(ResNet, self).__init__()
        self.stem=torch.nn.Sequential(
            torch.nn.Conv2d(3,32,kernel_size=(7,7),stride=(2,2),padding=1),
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=(3,3),stride=(2,2),padding=1),
        )
        self.layers1=self.build_cellblock(32,64,num_blocks[0],stride=2)
        self.layers2=self.build_cellblock(64,128,num_blocks[1],stride=2)
        self.layers3=self.build_cellblock(128,256,num_blocks[2],stride=2)
        self.layers4=self.build_cellblock(256,512,num_blocks[3],stride=2)

        self.avgpool = torch.nn.AdaptiveAvgPool2d(1)
        self.fc=torch.nn.Flatten()
        self.fc1=torch.nn.Linear(1*1*512,nb_class)

    def forward(self,x):
        x=self.stem(x)

        x=self.layers1(x)
        x=self.layers2(x)
        x=self.layers3(x)
        x=self.layers4(x)

        x=self.avgpool(x)
        x=self.fc(x)
        x=self.fc1(x)
        return x
    
    def build_cellblock(self,ch_in,ch_out,block,stride=1):
        rell_block=torch.nn.Sequential()
        rell_block.add_module('1',ConvBLK(ch_in,ch_out,stride))
        for _ in range(1,block):
            rell_block.add_module('2',ConvBLK(ch_out,ch_out,stride=1))
        return rell_block

# 创建resNet处理
def build_ResNet(ResName,nb_class):
    ResNet_Config={'ResNet18':[2,2,2,2]}
    return ResNet(ResNet_Config[ResName],nb_class)

if __name__ == '__main__':

    model=build_ResNet('ResNet18',3).cuda()
    criterion=torch.nn.CrossEntropyLoss().cuda()
    optimizer=torch.optim.Adam(model.parameters(),lr=0.01)
    for step in range(100):
        optimizer.zero_grad()

        h=model(train_x)
        cost=criterion(h,train_y.long())
        cost.backward()
        optimizer.step()
        pred = torch.argmax(model(train_x), 1)
        acc = (train_y == pred).float().mean()
        print(step,'loss={:.4f}'.format(cost.item()),'acc={:.2f}'.format(acc))

    pred=torch.argmax(model(train_x),1)
    acc=(train_y==pred).float().mean()
    print('train_acc',acc)

    pred = torch.argmax(model(test_x), 1)
    acc = (test_y == pred).float().mean()
    print('test_acc', acc)

